import os

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig


def download_qwen_model(model_name, save_dir):
    """下载Qwen模型（存在则跳过）"""
    # 1. 设置国内镜像，加速下载
    os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
    print(f"正在检查模型 {model_name} 是否已下载...")

    # 2. 检查本地目录是否存在，不存在则创建并下载
    if not os.path.exists(save_dir):
        os.makedirs(save_dir, exist_ok=True)
        print(f"模型未找到，开始下载至 {save_dir}...")
        # 执行下载命令，--resume-download 支持断点续传
        download_cmd = f"huggingface-cli download --resume-download {model_name} --local-dir {save_dir}"
        os.system(download_cmd)
        print("模型下载完成！")
    else:
        print("模型已存在，跳过下载步骤。")


def load_qwen_model(model_dir):
    """加载Qwen模型和Tokenizer（适配小模型的轻量化配置）"""
    print(f"正在加载模型 {model_dir}...")

    # 1. 加载Tokenizer（Qwen模型需指定trust_remote_code=True）
    tokenizer = AutoTokenizer.from_pretrained(
        model_dir,
        trust_remote_code=True,
        padding_side="right"  # 避免生成时警告
    )

    # 2. 配置4位量化（减少显存占用，适合0.6B小模型）
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )

    # 3. 加载模型（自动适配CPU/GPU）
    model = AutoModelForCausalLM.from_pretrained(
        model_dir,
        trust_remote_code=True,
        quantization_config=bnb_config,
        device_map="auto",  # 自动分配设备
        torch_dtype=torch.bfloat16
    )

    print("模型加载完成，可开始聊天（输入'退出'结束对话）！")
    return model, tokenizer


def chat_with_qwen(model, tokenizer):
    """交互式聊天循环"""
    # 初始化对话历史（支持多轮上下文）
    conversation = []

    while True:
        # 1. 获取用户输入
        user_input = input("\n你：")
        if user_input.strip() in ["退出", "quit", "exit"]:
            print("Qwen：再见！")
            break
        if not user_input.strip():
            print("Qwen：请输入有效内容哦～")
            continue

        # 2. 构建对话prompt（遵循Qwen的对话格式）
        conversation.append({"role": "user", "content": user_input})
        prompt = tokenizer.apply_chat_template(
            conversation,
            add_generation_prompt=True,  # 自动添加模型回复的前缀
            tokenize=False
        )

        # 3. 生成模型回复
        inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,  # 最大回复长度
            temperature=0.7,  # 回复随机性（0-1，越高越灵活）
            do_sample=True,  # 启用采样生成
            top_p=0.9,  # 核采样阈值
            pad_token_id=tokenizer.eos_token_id  # 避免padding警告
        )

        # 4. 解码并显示回复
        response = tokenizer.decode(
            outputs[0][len(inputs["input_ids"][0]):],  # 只取模型生成的部分
            skip_special_tokens=True
        )
        print(f"Qwen：{response}")

        # 5. 更新对话历史（保存上下文）
        conversation.append({"role": "assistant", "content": response})


if __name__ == "__main__":
    # 模型配置（可根据需求修改）
    MODEL_NAME = "Qwen/Qwen3-0.6B"
    SAVE_DIR = "./huggingface_cache"

    # 安装依赖（首次运行需执行，之后可注释）
    os.system("pip install -q transformers bitsandbytes torch accelerate")

    # 执行流程：下载 → 加载 → 聊天
    download_qwen_model(MODEL_NAME, SAVE_DIR)
    model, tokenizer = load_qwen_model(SAVE_DIR)
    chat_with_qwen(model, tokenizer)